import os
import operator
import asyncio
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

# --- 1. 环境与工具准备 (同教程8) ---
load_dotenv()
search_tool = TavilySearchResults(max_results=2)
all_tools = [search_tool]

# --- 2. 状态定义 ---
# 移除了 iteration_count，因为我们的控制将更加动态
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], operator.add]

# --- 3. 异步节点定义 ---
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))
model_with_tools = llm.bind_tools(all_tools)

# a. Agent & Router 节点 (简化，因为本教程重点是交互)
async def agent_node(state: AgentState):
    print("\n---AGENT: 思考中...---")
    response = await model_with_tools.ainvoke(state["messages"])
    return {"messages": [response]}

tool_node = ToolNode(all_tools)

def router_node(state: AgentState) -> str:
    last_message = state["messages"][-1]
    if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
        return END
    return "tools"

# --- 4. 构建图 ---
workflow = StateGraph(AgentState)
workflow.add_node("agent", agent_node)
workflow.add_node("tools", tool_node)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", router_node, {"tools": "tools", END: END})
workflow.add_edge("tools", "agent")

# --- 5. 编译图，所有工具都需要审批 ---
checkpointer = MemorySaver()
app = workflow.compile(checkpointer=checkpointer, interrupt_before=["tools"])

# --- 6. 核心：带时间旅行和修正功能的交互循环 ---
async def main():
    question = "德国的首都是哪里？那里的天气怎么样？"
    config = {"configurable": {"thread_id": "user_session_9"}}
    inputs = {"messages": [HumanMessage(content=question)]}

    while True:
        # 使用 astream_events 运行图直到中断
        async for event in app.astream_events(inputs, config=config, version="v2"):
            kind = event["event"]
            if kind == "on_chat_model_stream":
                content = event["data"]["chunk"].content
                if content:
                    print(content, end="", flush=True)

        inputs = None # 恢复运行时，输入应为 None
        
        # 获取当前状态快照
        snapshot = await app.aget_state(config)
        
        if not snapshot.next:
            print("\n\n---流程已完成！---")
            final_message = snapshot.values["messages"][-1]
            print("最终答案:")
            final_message.pretty_print()
            break
        
        print("\n\n---流程暂停，等待您的操作---")
        last_message = snapshot.values["messages"][-1]
        print("代理计划调用工具:", last_message.tool_calls)
        
        user_choice = input("您希望如何操作？ (yes/no/correct/exit): ")
        
        if user_choice.lower() == "yes":
            print("\n---用户批准，继续执行---")
            # `inputs` 已经为 None，所以循环会从中断点恢复
            continue
        elif user_choice.lower() == "no" or user_choice.lower() == "exit":
            print("\n---用户否决或退出，流程终止---")
            break
        elif user_choice.lower() == "correct":
            correction = input("请输入您的修正指令: ")
            
            # 关键：获取历史状态
            history = [s async for s in app.aget_state_history(config)]
            
            # 定位到代理做出错误决定之前的那个状态
            last_agent_state_snapshot = history[-2]
            
            # 创建一个 HumanMessage 来直接、清晰地给出修正指令
            # 这比伪造 ToolMessage 更健壮，也更符合模型的对话模式
            correction_message = HumanMessage(
                content=correction,
            )
            
            # 使用 `aupdate_state` 从历史快照点开始，强行插入我们的修正信息
            await app.aupdate_state(
                last_agent_state_snapshot.config,
                {"messages": [correction_message]},
            )
            print("\n---状态已修正，代理将从上一步重新思考---")
            continue
        else:
            print(f"无法识别的指令 '{user_choice}'，流程终止。")
            break

if __name__ == "__main__":
    asyncio.run(main()) 